package com.bear.atlassianmcp.common.util;

import io.modelcontextprotocol.server.McpAsyncServerExchange;
import io.modelcontextprotocol.server.McpSyncServerExchange;
import io.modelcontextprotocol.server.transport.WebFluxSseServerTransportProvider;
import io.modelcontextprotocol.spec.McpServerSession;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.mcp.McpToolUtils;

import java.lang.reflect.Field;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @Auther: bear
 * @Date: 2025/08/18/11:00
 * @Description:
 */
public class SessionUtil {
    static final Field exchangeField;
    static final Field sessionField;
    static final Field idField;
    static final Field transportField;
    static final Field session2headersField;

    static {
        exchangeField = getField(McpSyncServerExchange.class, "exchange");
        sessionField = getField(McpAsyncServerExchange.class, "session");
        idField = getField(McpServerSession.class, "id");
        transportField = getField(McpServerSession.class, "transport");
        session2headersField = getField(WebFluxSseServerTransportProvider.class, "session2headers");
    }

    static Field getField(Class<?> clazz, String fieldName) {
        try {
            Field field = clazz.getDeclaredField(fieldName);
            field.setAccessible(true);
            return field;
        } catch (Exception e) {
            return null;
        }
    }

    public static String getSessionId(ToolContext toolContext) {
        try {
            Optional<McpSyncServerExchange> mcpExchange = McpToolUtils.getMcpExchange(toolContext);
            McpSyncServerExchange syncServerExchange = mcpExchange.get();
            McpAsyncServerExchange asyncServerExchange = (McpAsyncServerExchange) exchangeField.get(syncServerExchange);
            McpServerSession session = (McpServerSession) sessionField.get(asyncServerExchange);
            return (String) idField.get(session);
        } catch (Exception e) {
            return null;
        }
    }

    /**
     * 获取session对应的HTTP请求头信息
     * @param toolContext 工具上下文
     * @return session对应的请求头Map，如果获取失败返回null
     */
    public static Map<String, String> getSessionHeaders(ToolContext toolContext) {
        try {
            Optional<McpSyncServerExchange> mcpExchange = McpToolUtils.getMcpExchange(toolContext);
            McpSyncServerExchange syncServerExchange = mcpExchange.get();
            McpAsyncServerExchange asyncServerExchange = (McpAsyncServerExchange) exchangeField.get(syncServerExchange);
            McpServerSession session = (McpServerSession) sessionField.get(asyncServerExchange);

            // 获取session的transport对象
            Object transport = transportField.get(session);

            // 通过transport获取WebFluxSseServerTransportProvider实例
            // 需要通过transport的外部类引用获取
            Field thisField = transport.getClass().getDeclaredField("this$0");
            thisField.setAccessible(true);
            WebFluxSseServerTransportProvider transportProvider = (WebFluxSseServerTransportProvider) thisField.get(transport);

            // 获取session2headers字段
            ConcurrentHashMap<String, Map<String, String>> session2headers =
                (ConcurrentHashMap<String, Map<String, String>>) session2headersField.get(transportProvider);

            // 获取当前session的ID
            String sessionId = (String) idField.get(session);

            // 返回对应的headers
            return session2headers.get(sessionId);
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    /**
     * 获取session对应的HTTP请求头中的Authorization信息
     * @param toolContext 工具上下文
     * @return session对应的Authorization信息，如果获取失败返回null
     */
    public static String getAuthorization(ToolContext toolContext) {
        Map<String, String> sessionHeaders = getSessionHeaders(toolContext);
        return sessionHeaders == null ? null : sessionHeaders.get("authorization");
    }
}
